package sc.digital.human;

import com.google.common.base.Splitter;
import com.google.gson.Gson;
import com.google.gson.JsonArray;
import com.google.gson.JsonObject;
import com.sun.net.httpserver.HttpExchange;
import com.sun.net.httpserver.HttpHandler;
import com.sun.net.httpserver.HttpsConfigurator;
import com.sun.net.httpserver.HttpsServer;
import okhttp3.*;
import okio.BufferedSink;
import okio.BufferedSource;
import okio.ByteString;
import okio.Okio;

import javax.annotation.Nonnull;
import javax.net.ssl.KeyManagerFactory;
import javax.net.ssl.SSLContext;
import javax.net.ssl.TrustManagerFactory;
import javax.swing.*;
import javax.swing.text.*;
import java.awt.*;
import java.awt.event.ActionListener;
import java.io.File;
import java.io.IOException;
import java.io.OutputStream;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.URI;
import java.net.UnknownHostException;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.security.KeyStore;
import java.security.SecureRandom;
import java.util.Base64;
import java.util.Map;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

public class Server extends JFrame implements HttpHandler {
    private static final String KET_STORE_NAME = "/123456.jks";
    private static String KET_STORE_FILE = "", KET_STORE_PASSWORD = "123456";
    private static int SERVER_PORT = 8080;
    private final Setup setup = new Setup(this);
    private final Gson gson = new Gson();
    private final Credentials credentials = new Credentials();
    private final OkHttpClient httpClient = new OkHttpClient.Builder().readTimeout(60, TimeUnit.SECONDS).build();
    private HttpsServer httpServer;

    public Server() {
        super.setTitle("三川数字人（服务）");
        super.setIconImage(Config.LOGO);
        super.setPreferredSize(new Dimension(800, 600));
        super.setDefaultCloseOperation(EXIT_ON_CLOSE);
        super.setLayout(new BorderLayout());
        super.add(new JToolBar() {{
            this.add(new JButton(httpServer == null ? "▶" : "■") {{
                this.setFocusable(false);
                this.setFocusPainted(false);
                this.addActionListener(server);
            }});
            this.add(new JComboBox<Browse>(Browse.values()) {{
                this.setFocusable(false);
                this.addActionListener(browse);
            }});
            this.add(new JButton("⚙") {{
                this.setFocusable(false);
                this.setFocusPainted(false);
                this.addActionListener((e) -> setup.setVisible(true));
            }});
        }}, BorderLayout.NORTH);
        this.add(new JScrollPane(new JTextPane() {{
            this.setEditable(false);
            this.setBackground(Color.BLACK);
            System.setOut(new PrintStream(this, Color.WHITE, System.out));
            System.setErr(new PrintStream(this, Color.RED, System.err));
        }}), BorderLayout.CENTER);
        super.pack();
        super.setLocationRelativeTo(null);
    }

    public void start() throws Exception {
        if (this.httpServer == null) {
            this.httpServer = HttpsServer.create(new InetSocketAddress(SERVER_PORT), 0);
            this.httpServer.createContext("/", this);
            this.httpServer.createContext("/api/asr", new AsrHandler());
            this.httpServer.createContext("/api/tts", new TtsHandler());
            this.httpServer.createContext("/api/chat", new ChatHandler());
            this.httpServer.createContext("/api/chat/wenxin", new WenxinHandler());
            KeyStore keyStore = KeyStore.getInstance(KeyStore.getDefaultType());
            keyStore.load(new File(KET_STORE_FILE).isFile()
                            ? Files.newInputStream(Paths.get(KET_STORE_FILE))
                            : getClass().getResourceAsStream(KET_STORE_NAME)
                    , KET_STORE_PASSWORD.toCharArray());
            KeyManagerFactory keyManager = KeyManagerFactory.getInstance(KeyManagerFactory.getDefaultAlgorithm());
            keyManager.init(keyStore, KET_STORE_PASSWORD.toCharArray());
            TrustManagerFactory trustManager = TrustManagerFactory.getInstance(TrustManagerFactory.getDefaultAlgorithm());
            trustManager.init(keyStore);
            SSLContext sslContext = SSLContext.getInstance("TLS");
            sslContext.init(keyManager.getKeyManagers(), trustManager.getTrustManagers(), new SecureRandom());
            this.httpServer.setHttpsConfigurator(new HttpsConfigurator(sslContext));
            this.httpServer.setExecutor(Executors.newCachedThreadPool());
            this.httpServer.start();
        }
    }

    public void stop() {
        this.httpServer.stop(0);
        this.httpServer = null;
    }

    void refresh(String appKey, String secretKey, Credentials credentials) throws IOException {
        if (credentials.access_token == null || credentials.access_token.isEmpty()
                || credentials.expires_in < System.currentTimeMillis()) {
            try (Response response = this.httpClient.newCall(new Request.Builder()
                    .url(String.format(Config.TOKEN_URL, appKey, secretKey))
                    .build()).execute()) {
                String json = response.body().string();
                System.out.printf("鉴权认证：%s%n", json);
                Credentials credential = this.gson.fromJson(json, Credentials.class);
                credentials.access_token = credential.access_token;
                credentials.expires_in = System.currentTimeMillis() + credential.expires_in * 1000;
            }
        }
    }

    JsonObject query(String sessionId, String query) {
        JsonObject questions = new JsonObject();
        questions.addProperty("version", Config.CHAT_VERSION);
        questions.addProperty("service_id", Config.SERVICE_ID);
        JsonArray skill_ids = new JsonArray();
        if (Config.SKILL_ID != null) {
            for (String skill_id : Config.SKILL_ID) {
                skill_ids.add(skill_id);
            }
        }
        questions.add("skill_ids", skill_ids);
        questions.addProperty("log_id", UUID.randomUUID().toString());
        questions.addProperty("session_id", sessionId);
        JsonObject request = new JsonObject();
        request.addProperty("terminal_id", "UNIT_DEV_" + Config.APP_ID);
        request.addProperty("user_id", "UNIT_DEV_" + Config.APP_ID);
        request.addProperty("query", query);
        JsonObject query_info = new JsonObject();
        query_info.addProperty("type", "TEXT");
        query_info.addProperty("source", "ASR");
        request.add("query_info", query_info);
        questions.add("request", request);
        return questions;
    }

    @Override
    public void handle(HttpExchange httpExchange) throws IOException {
        System.out.printf("URL：%s%n", httpExchange.getRequestURI());
        try (BufferedSource source = Okio.buffer(Okio.source(Objects.requireNonNull(getClass()
                .getResourceAsStream(httpExchange.getRequestURI().getPath()))))) {
            ByteString byteString = source.readByteString();
            httpExchange.sendResponseHeaders(200, byteString.size());
            try (BufferedSink sink = Okio.buffer(Okio.sink(httpExchange.getResponseBody()))) {
                sink.write(byteString);
            }
        }
        httpExchange.close();
    }

    ActionListener server = e -> {
        ((JButton) e.getSource()).setEnabled(false);
        new Thread(() -> {
            if (this.httpServer == null) {
                try {
                    this.start();
                } catch (Exception ex) {
                    JOptionPane.showMessageDialog(this, ex.toString(), ex.getMessage(), JOptionPane.PLAIN_MESSAGE);
                }
            } else {
                this.stop();
            }
            SwingUtilities.invokeLater(() -> {
                ((JButton) e.getSource()).setEnabled(true);
                ((JButton) e.getSource()).setText(this.httpServer == null ? "▶" : "■");
            });
        }).start();
    };
    ActionListener browse = e -> {
        if (this.httpServer != null) {
            new Thread(() -> {
                try {
                    Desktop.getDesktop().browse(new URI(((Browse) Objects.requireNonNull(((JComboBox<?>) e.getSource()).getSelectedItem())).toURL()));
                } catch (Exception ex) {
                    JOptionPane.showMessageDialog(this, ex.toString(), ex.getMessage(), JOptionPane.PLAIN_MESSAGE);
                }
            }).start();
        }
    };

    class AsrHandler implements HttpHandler {
        @Override
        public void handle(HttpExchange httpExchange) throws IOException {
            System.out.printf("API：%s%n", httpExchange.getRequestURI());
            refresh(Config.APP_KEY, Config.SECRET_KEY, credentials);
            try (BufferedSource source = Okio.buffer(Okio.source(httpExchange.getRequestBody()))) {
                try (Response response = httpClient.newCall(new Request.Builder()
                        .header("format", "pcm")
                        .header("rate", String.valueOf(Config.SAMPLE_RATE))
                        .post(RequestBody.create(MediaType.parse("audio/pcm;rate=" + Config.SAMPLE_RATE)
                                , Base64.getDecoder().decode(source.readByteString().toByteArray())))
                        .url(String.format(Config.ASR_URL, credentials.access_token))
                        .build()).execute()) {
                    ByteString byteString = response.body().source().readByteString();
                    System.out.printf("语音识别：%s%n", byteString.utf8());
                    httpExchange.getResponseHeaders().set("Content-Type", response.header("Content-Type"));
                    httpExchange.sendResponseHeaders(response.code(), byteString.size());
                    try (BufferedSink sink = Okio.buffer(Okio.sink(httpExchange.getResponseBody()))) {
                        sink.write(byteString);
                    }
                }
            }
            httpExchange.close();
        }
    }

    class TtsHandler implements HttpHandler {
        @Override
        public void handle(HttpExchange httpExchange) throws IOException {
            System.out.printf("API：%s%n", httpExchange.getRequestURI());
            refresh(Config.APP_KEY, Config.SECRET_KEY, credentials);
            Map<String, String> map = Splitter.on("&").withKeyValueSeparator("=")
                    .split(httpExchange.getRequestURI().getQuery());
            try (Response response = httpClient.newCall(new Request.Builder()
                    .url(String.format(Config.TTS_URL
                            , Tsn.urlEncoder(map.get("text")), Tsn.Per.PER0.key, Tsn.Aue.MP3.key, credentials.access_token)).build()).execute()) {
                ByteString byteString = response.body().source().readByteString();
                System.out.printf("语音合成：%s%n", byteString);
                httpExchange.getResponseHeaders().set("Content-Type", response.header("Content-Type"));
                httpExchange.sendResponseHeaders(response.code(), byteString.size());
                try (BufferedSink sink = Okio.buffer(Okio.sink(httpExchange.getResponseBody()))) {
                    sink.write(byteString);
                }
            }
            httpExchange.close();
        }
    }

    class ChatHandler implements HttpHandler {
        @Override
        public void handle(HttpExchange httpExchange) throws IOException {
            System.out.printf("API：%s%n", httpExchange.getRequestURI());
            refresh(Config.APP_KEY, Config.SECRET_KEY, credentials);
            Map<String, String> map = Splitter.on("&").withKeyValueSeparator("=")
                    .split(httpExchange.getRequestURI().getQuery());
            try (Response response = httpClient.newCall(new Request.Builder()
                    .post(RequestBody.create(MediaType.parse("application/json")
                            , gson.toJson(query(map.get("session_id"), map.get("query")))))
                    .url(String.format(Chat.Ver.toURL(Config.CHAT_VERSION), credentials.access_token))
                    .build()).execute()) {
                ByteString byteString = response.body().source().readByteString();
                System.out.printf("智能对话：%s%n", byteString.utf8());
                httpExchange.getResponseHeaders().set("Content-Type", response.header("Content-Type"));
                httpExchange.sendResponseHeaders(response.code(), byteString.size());
                try (BufferedSink sink = Okio.buffer(Okio.sink(httpExchange.getResponseBody()))) {
                    sink.write(byteString);
                }
            }
            httpExchange.close();
        }
    }

    class WenxinHandler implements HttpHandler {
        Credentials credentials = new Credentials();

        @Override
        public void handle(HttpExchange httpExchange) throws IOException {
            System.out.printf("API：%s%n", httpExchange.getRequestURI());
            refresh(Config.CHAT_APP_KEY, Config.CHAT_SECRET_KEY, credentials);
            try (BufferedSource source = Okio.buffer(Okio.source(httpExchange.getRequestBody()))) {
                ByteString byteString = source.readByteString();
                System.out.printf("问：%s%n", byteString.utf8());
                try (Response response = httpClient.newCall(new Request.Builder()
                        .post(RequestBody.create(MediaType.parse("application/json"), byteString))
                        .url(String.format(Chat.Bot.toURL(Config.CHAT_BOT_NAME), credentials.access_token))
                        .build()).execute()) {
                    byteString = response.body().source().readByteString();
                    System.out.printf("答：%s%n", byteString.utf8());
                    httpExchange.getResponseHeaders().set("Content-Type", response.header("Content-Type"));
                    httpExchange.sendResponseHeaders(response.code(), byteString.size());
                    try (BufferedSink sink = Okio.buffer(Okio.sink(httpExchange.getResponseBody()))) {
                        sink.write(byteString);
                    }
                }
            }
            httpExchange.close();
        }
    }

    static class Credentials {
        public String access_token;
        public long expires_in;
    }

    enum Browse {
        UNIT("智能对话", "chat/index.html"), ERNIE("文心一言", "chat/wenxin.html");
        final String name;
        final String url;

        Browse(String name, String url) {
            this.name = name;
            this.url = url;
        }

        public String toURL() {
            try {
                return String.format("https://%s:%s/%s", InetAddress.getLocalHost().getHostAddress(), SERVER_PORT, url);
            } catch (UnknownHostException e) {
                return String.format("https://%s:%s/%s", "127.0.0.1", SERVER_PORT, url);
            }
        }

        @Override
        public String toString() {
            return String.format("%s | %s", name, toURL());
        }
    }

    static class PrintStream extends java.io.PrintStream {
        static int ELEMENT_COUNT = 5000;
        final Document document;
        final Element element;
        final SimpleAttributeSet attr;

        public PrintStream(JTextPane logger, Color color, OutputStream out) {
            super(out, true);
            this.document = logger.getDocument();
            this.element = document.getDefaultRootElement();
            StyleConstants.setForeground(this.attr = new SimpleAttributeSet(), color);
        }

        @Override
        public void write(int b) {
            this.write(new byte[]{(byte) b}, 0, 1);
        }

        @Override
        public void write(@Nonnull byte[] b) {
            this.write(b, 0, b.length);
        }

        @Override
        public void write(@Nonnull byte[] buffer, int offset, int length) {
            String doc = new String(buffer, offset, length);
            SwingUtilities.invokeLater(() -> {
                try {
                    if (this.element.getElementCount() > ELEMENT_COUNT) {
                        this.document.remove(this.element.getElement(0).getStartOffset(), this.element.getElement(0).getEndOffset());
                    }
                    this.document.insertString(this.document.getLength(), doc, this.attr);
                } catch (BadLocationException e) {
                    throw new RuntimeException(e);
                }
            });
        }
    }

    static void readCmdlineArgs(String[] args) {
        KET_STORE_FILE = Config.getCmdline(args, "key.store.file", KET_STORE_FILE);
        if (new File(KET_STORE_FILE).isFile()) {
            KET_STORE_PASSWORD = Config.getCmdline(args, "key.store.password", KET_STORE_PASSWORD);
        }
        SERVER_PORT = Config.getCmdline(args, "server.port", SERVER_PORT);
        System.out.println("Usage: [--key.store.file=证书文件 --key.store.password=证书密码 --server.port=服务端口]");
        System.out.println("--key.store.file              证书文件：" + (new File(KET_STORE_FILE).isFile() ? KET_STORE_FILE : KET_STORE_NAME));
        System.out.println("--key.store.password          证书密码：" + KET_STORE_PASSWORD);
        System.out.println("--server.port                 服务端口：" + SERVER_PORT);
    }

    public static void main(String[] args) throws Exception {
        Config.readConfig();
        Config.readCmdlineArgs(args);
        readCmdlineArgs(args);
        UIManager.setLookAndFeel(UIManager.getSystemLookAndFeelClassName());
        SwingUtilities.invokeLater(() -> new Server().setVisible(true));
    }
}
